{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ecfdbe35-92d2-4151-86a4-ae054d6d2344",
   "metadata": {},
   "source": [
    "# AlloyDB Data Prep Notebook\n",
    "\n",
    "This notebook downloads the dataset [\"Financial Transactions Dataset: Analytics\"](https://www.kaggle.com/datasets/computingvictor/transactions-fraud-datasets) from Kaggle and runs various transformations to import the data into an AlloyDB for PostgreSQL database. \n",
    "\n",
    "See the following documentation for more details on the steps performed in the notebook: [Import a CSV File to AlloyDB](https://cloud.google.com/alloydb/docs/import-csv-file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec723d99-da14-467a-9111-895b304b80f0",
   "metadata": {},
   "source": [
    "## Basic Setup"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dabae54e-f230-44d8-8e4c-e69478d7c0f5",
   "metadata": {},
   "source": [
    "### Define Notebook Variables\n",
    "\n",
    "Update the variables below to match your environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aea3732-c4c2-40fa-9cc5-98114d48b7e7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "project_id = \"my-project\"\n",
    "region = \"my-region\"\n",
    "alloydb_cluster = \"my-alloydb-cluster\"\n",
    "alloydb_instance = \"my-alloydb-instance\"\n",
    "alloydb_database = \"finance\"\n",
    "alloydb_password = input(\"Please provide a password to be used for 'postgres' database user: \")\n",
    "gcs_bucket_name = f\"project-files-{project_id}\"\n",
    "export_staging_directory = f\"gs://{gcs_bucket_name}/staging\"\n",
    "export_output_directory = f\"gs://{gcs_bucket_name}/output\"\n",
    "vpc = \"demo-vpc\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b212dc5f-ff39-4066-9a3d-d3749ce7dc81",
   "metadata": {},
   "source": [
    "### Connect to your Google Cloud Project"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73121f7a-fab3-49e6-b764-5af406e2e9db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Configure gcloud.\n",
    "!gcloud config set project {project_id}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04fe21f1-7040-42c4-bf28-8938c2797bd2",
   "metadata": {},
   "source": [
    "### Configure Logging"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98044a4f-d76a-4331-b093-d6a2ace18016",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "# Configure the root logger to output messages with INFO level or above\n",
    "logging.basicConfig(level=logging.INFO, stream=sys.stdout, format='%(asctime)s[%(levelname)5s][%(name)14s] - %(message)s',  datefmt='%H:%M:%S', force=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acb23dd4-0818-4cf1-8fac-6ec7d61ca6e1",
   "metadata": {},
   "source": [
    "### Install Dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59c21d20-72f8-4b13-9426-4f19ccaa2c3b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! pip install --quiet google-cloud-storage==2.19.0 \\\n",
    "                      google-cloud-aiplatform==1.74.0 \\\n",
    "                      asyncpg==0.30.0 \\\n",
    "                      google.cloud.alloydb.connector==1.9.0 \\\n",
    "                      google-genai==1.4.0 \\\n",
    "                      kagglehub==0.3.12"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c2a2bbe-da57-40b6-8fa9-5b4359355ea3",
   "metadata": {},
   "source": [
    "### Define Helper Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c082c9-1c8f-418f-b6dc-576d26afa085",
   "metadata": {},
   "source": [
    "#### rest_api_helper()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ee81d6-be7f-44aa-a4e7-95d251b65197",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "import google.auth\n",
    "import json\n",
    "\n",
    "# Get an access token based upon the current user\n",
    "creds, _ = google.auth.default()\n",
    "authed_session = google.auth.transport.requests.AuthorizedSession(creds)\n",
    "access_token=creds.token\n",
    "\n",
    "if project_id:\n",
    "  authed_session.headers.update({\"x-goog-user-project\": project_id}) # Required to workaround a project quota bug\n",
    "\n",
    "def rest_api_helper(\n",
    "    session: requests.Session,\n",
    "    url: str,\n",
    "    http_verb: str,\n",
    "    request_body: dict = None,\n",
    "    params: dict = None\n",
    "  ) -> dict:\n",
    "  \"\"\"Calls a REST API using a pre-authenticated requests Session.\"\"\"\n",
    "\n",
    "  headers = {\"Content-Type\": \"application/json\"}\n",
    "\n",
    "  try:\n",
    "\n",
    "    if http_verb == \"GET\":\n",
    "      response = session.get(url, headers=headers, params=params)\n",
    "    elif http_verb == \"POST\":\n",
    "      response = session.post(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"PUT\":\n",
    "      response = session.put(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"PATCH\":\n",
    "      response = session.patch(url, json=request_body, headers=headers, params=params)\n",
    "    elif http_verb == \"DELETE\":\n",
    "      response = session.delete(url, headers=headers, params=params)\n",
    "    else:\n",
    "      raise ValueError(f\"Unknown HTTP verb: {http_verb}\")\n",
    "\n",
    "    # Raise an exception for bad status codes (4xx or 5xx)\n",
    "    response.raise_for_status()\n",
    "\n",
    "    # Check if response has content before trying to parse JSON\n",
    "    if response.content:\n",
    "        return response.json()\n",
    "    else:\n",
    "        return {} # Return empty dict for empty responses (like 204 No Content)\n",
    "\n",
    "  except requests.exceptions.RequestException as e:\n",
    "      # Catch potential requests library errors (network, timeout, etc.)\n",
    "      # Log detailed error information\n",
    "      print(f\"Request failed: {e}\")\n",
    "      if e.response is not None:\n",
    "          print(f\"Request URL: {e.request.url}\")\n",
    "          print(f\"Request Headers: {e.request.headers}\")\n",
    "          print(f\"Request Body: {e.request.body}\")\n",
    "          print(f\"Response Status: {e.response.status_code}\")\n",
    "          print(f\"Response Text: {e.response.text}\")\n",
    "          # Re-raise a more specific error or a custom one\n",
    "          raise RuntimeError(f\"API call failed with status {e.response.status_code}: {e.response.text}\") from e\n",
    "      else:\n",
    "          raise RuntimeError(f\"API call failed: {e}\") from e\n",
    "  except json.JSONDecodeError as e:\n",
    "      print(f\"Failed to decode JSON response: {e}\")\n",
    "      print(f\"Response Text: {response.text}\")\n",
    "      raise RuntimeError(f\"Invalid JSON received from API: {response.text}\") from e\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c6840fa-2d3e-48a1-bf36-93f66d46f1b0",
   "metadata": {},
   "source": [
    "#### run_query()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f146496d-be74-4c68-972d-59998847819e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Create AlloyDB Query Helper Function\n",
    "import sqlalchemy\n",
    "from sqlalchemy import text, exc\n",
    "import pandas as pd\n",
    "\n",
    "async def run_query(pool, sql: str, params = None, output_as_df: bool = True):\n",
    "    \"\"\"Executes a SQL query or statement against the database pool.\n",
    "\n",
    "    Handles various SQL statements:\n",
    "    - SELECT/WITH: Returns results as a DataFrame (if output_as_df=True)\n",
    "      or ResultProxy. Supports parameters. Does not commit.\n",
    "    - EXPLAIN/EXPLAIN ANALYZE: Executes the explain, returns the query plan\n",
    "      as a formatted multi-line string. Ignores output_as_df.\n",
    "      Supports parameters. Does not commit.\n",
    "    - INSERT/UPDATE/DELETE/CREATE/ALTER etc.: Executes the statement,\n",
    "      commits the transaction, logs info, and returns the ResultProxy.\n",
    "      Supports single or bulk parameters (executemany).\n",
    "\n",
    "    Args:\n",
    "      pool: An asynchronous SQLAlchemy connection pool.\n",
    "      sql: A string containing the SQL query or statement template.\n",
    "      params: Optional.\n",
    "        - None: Execute raw SQL (Use with caution for non-SELECT/EXPLAIN).\n",
    "        - dict or tuple: Parameters for a single execution.\n",
    "        - list of dicts/tuples: Parameters for bulk execution (executemany).\n",
    "      output_as_df (bool): If True and query is SELECT/WITH, return pandas DataFrame.\n",
    "                           Ignored for EXPLAIN and non-data-returning statements.\n",
    "\n",
    "    Returns:\n",
    "      pandas.DataFrame | str | sqlalchemy.engine.Result | None:\n",
    "        - DataFrame: For SELECT/WITH if output_as_df=True.\n",
    "        - str: For EXPLAIN/EXPLAIN ANALYZE, containing the formatted query plan.\n",
    "        - ResultProxy: For non-SELECT/WITH/EXPLAIN statements, or SELECT/WITH\n",
    "                       if output_as_df=False.\n",
    "        - None: If a SQLAlchemy ProgrammingError or other specific error occurs.\n",
    "\n",
    "    Raises:\n",
    "        Exception: Catches and logs `sqlalchemy.exc.ProgrammingError`, returning None.\n",
    "                   May re-raise other database exceptions.\n",
    "\n",
    "    Example Execution:\n",
    "      Single SELECT:\n",
    "        sql_select = \"SELECT ticker, company_name from investments LIMIT 5\"\n",
    "        df_result = await run_query(pool, sql_select)\n",
    "\n",
    "      Single non-SELECT - Parameterized (Safe!):\n",
    "        Parameterized INSERT:\n",
    "          sql_insert = \"INSERT INTO investments (ticker, company_name) VALUES (:ticker, :name)\"\n",
    "          params_insert = {\"ticker\": \"NEW\", \"name\": \"New Company\"}\n",
    "          insert_result = await run_query(pool, sql_insert, params_insert)\n",
    "\n",
    "        Parameterized UPDATE:\n",
    "          sql_update = \"UPDATE products SET price = :price WHERE id = :product_id\"\n",
    "          params_update = {\"price\": 99.99, \"product_id\": 123}\n",
    "          update_result = await run_query(pool, sql_update, params_update)\n",
    "\n",
    "      Bulk Update:\n",
    "        docs = pd.DataFrame([\n",
    "            {'id': 101, 'sparse_embedding': '[0.1, 0.2]'},\n",
    "            {'id': 102, 'sparse_embedding': '[0.3, 0.4]'},\n",
    "            # ... more rows\n",
    "        ])\n",
    "\n",
    "        update_sql_template = '''\n",
    "            UPDATE products\n",
    "            SET sparse_embedding = :embedding,\n",
    "                sparse_embedding_model = 'BM25'\n",
    "            WHERE id = :product_id\n",
    "        ''' # Using named parameters :param_name\n",
    "\n",
    "        # Prepare list of dictionaries for params\n",
    "        data_to_update = [\n",
    "            {\"embedding\": row.sparse_embedding, \"product_id\": row.id}\n",
    "            for row in docs.itertuples(index=False)\n",
    "        ]\n",
    "\n",
    "        if data_to_update:\n",
    "          bulk_result = await run_query(pool, update_sql_template, data_to_update)\n",
    "          # bulk_result is the SQLAlchemy ResultProxy\n",
    "\n",
    "    \"\"\"\n",
    "    sql_lower_stripped = sql.strip().lower()\n",
    "    is_select_with = sql_lower_stripped.startswith(('select', 'with'))\n",
    "    is_explain = sql_lower_stripped.startswith('explain')\n",
    "\n",
    "    # Determine if the statement is expected to return data rows or a plan\n",
    "    is_data_returning = is_select_with or is_explain\n",
    "\n",
    "    # Determine actual DataFrame output eligibility (only for SELECT/WITH)\n",
    "    effective_output_as_df = output_as_df and is_select_with\n",
    "\n",
    "    # Check if params suggest a bulk operation (for logging purposes)\n",
    "    is_bulk_operation = isinstance(params, (list, tuple)) and len(params) > 0 and isinstance(params[0], (dict, tuple, list))\n",
    "\n",
    "    async with pool.connect() as conn:\n",
    "        try:\n",
    "          # Execute with or without params\n",
    "          if params:\n",
    "              result = await conn.execute(text(sql), params)\n",
    "          else:\n",
    "              # Add warning for raw SQL only if it's NOT data-returning\n",
    "              #if not is_data_returning:\n",
    "                  #logging.warning(\"Executing non-SELECT/EXPLAIN raw SQL without parameters. Ensure SQL is safe.\")\n",
    "              result = await conn.execute(text(sql))\n",
    "\n",
    "          # --- Handle statements that return data or plan ---\n",
    "          if is_data_returning:\n",
    "              if is_explain:\n",
    "                  # Fetch and format EXPLAIN output as a string\n",
    "                    try:\n",
    "                        plan_rows = result.fetchall()\n",
    "                        # EXPLAIN output is usually text in the first column\n",
    "                        query_plan = \"\\n\".join([str(row[0]) for row in plan_rows])\n",
    "                        #logging.info(f\"EXPLAIN executed successfully for: {sql[:100]}...\")\n",
    "                        return query_plan\n",
    "                    except Exception as e:\n",
    "                        logging.error(f\"Error fetching/formatting EXPLAIN result: {e}\")\n",
    "                        return None\n",
    "              else: # Handle SELECT / WITH\n",
    "                  if effective_output_as_df:\n",
    "                      try:\n",
    "                          rows = result.fetchall()\n",
    "                          column_names = result.keys()\n",
    "                          df = pd.DataFrame(rows, columns=column_names)\n",
    "                          #logging.info(f\"SELECT/WITH executed successfully, returning DataFrame for: {sql[:100]}...\")\n",
    "                          return df\n",
    "                      except Exception as e:\n",
    "                          logging.error(f\"Error converting SELECT result to DataFrame: {e}\")\n",
    "                          logging.info(f\"Returning raw ResultProxy for SELECT/WITH due to DataFrame conversion error for: {sql[:100]}...\")\n",
    "                          return result # Fallback to raw result\n",
    "                  else:\n",
    "                      # Return raw result proxy for SELECT/WITH if df output not requested\n",
    "                      #logging.info(f\"SELECT/WITH executed successfully, returning ResultProxy for: {sql[:100]}...\")\n",
    "                      return result\n",
    "\n",
    "          # --- Handle Non-Data Returning Statements (INSERT, UPDATE, DELETE, CREATE, etc.) ---\n",
    "          else:\n",
    "              await conn.commit() # Commit changes ONLY for these statements\n",
    "              operation_type = sql.strip().split()[0].upper()\n",
    "              row_count = result.rowcount # Note: rowcount behavior varies\n",
    "\n",
    "              if is_bulk_operation:\n",
    "                  print(f\"Bulk {operation_type} executed for {len(params)} items. Result rowcount: {row_count}\")\n",
    "              elif operation_type in ['INSERT', 'UPDATE', 'DELETE']:\n",
    "                  print(f\"{operation_type} statement executed successfully. {row_count} row(s) affected.\")\n",
    "              else: # CREATE, ALTER, etc.\n",
    "                  print(f\"{operation_type} statement executed successfully. Result rowcount: {row_count}\")\n",
    "              return result # Return the result proxy\n",
    "\n",
    "        except exc.ProgrammingError as e:\n",
    "            # Log the error with context\n",
    "            logging.error(f\"SQL Programming Error executing query:\\nSQL: {sql[:500]}...\\nParams (sample): {str(params)[:500]}...\\nError: {e}\")\n",
    "            # Rollback might happen automatically on context exit with error, but explicit can be clearer\n",
    "            # await conn.rollback() # Consider if needed based on pool/transaction settings\n",
    "            return None # Return None on handled programming errors\n",
    "        except Exception as e:\n",
    "            # Log other unexpected errors\n",
    "            logging.error(f\"An unexpected error occurred during query execution:\\nSQL: {sql[:500]}...\\nError: {e}\")\n",
    "            # await conn.rollback() # Consider if needed\n",
    "            raise # Re-raise unexpected errors\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f764979b-3e3f-409d-94f5-09b867e068f2",
   "metadata": {},
   "source": [
    "## Prepare Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bac2267b-7bad-4730-b56b-6ff87c2ecbf1",
   "metadata": {},
   "source": [
    "### Download Kaggle Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbbc8c2f-e55c-40f4-84b0-9055feabcc0e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import kagglehub\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "# Set the handle for the dataset you want to download\n",
    "dataset = \"computingvictor/transactions-fraud-datasets\"\n",
    "\n",
    "dataset_files = [\"transactions_data.csv\", \"cards_data.csv\", \"users_data.csv\", \"mcc_codes.json\", \"train_fraud_labels.json\"]\n",
    "\n",
    "os.makedirs('../data/', exist_ok=True) \n",
    "\n",
    "for file_to_download in dataset_files:\n",
    "    # Set the local path where you want to save the file\n",
    "    # In this case, it will be saved in the ../data/ directory\n",
    "    local_file_path = '../data/' + file_to_download\n",
    "\n",
    "    # Download the specified file from the dataset\n",
    "    # The function returns the path to the downloaded file\n",
    "    downloaded_path = kagglehub.dataset_download(\n",
    "        dataset,\n",
    "        path=file_to_download,\n",
    "    )\n",
    "\n",
    "    # The file is downloaded to a cache directory managed by kagglehub.\n",
    "    # This part of the code will move it to your desired local path.\n",
    "    # If a file with the same name already exists at the destination, it will be overwritten.\n",
    "    shutil.move(downloaded_path, local_file_path)\n",
    "\n",
    "    print(f\"File '{file_to_download}' downloaded to: {os.path.abspath(local_file_path)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5f42350-e11a-4ef5-8e5b-17c3480ff5c6",
   "metadata": {},
   "source": [
    "### Reformat JSON files as CSV"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56ff251a-1702-41f5-8fd9-380da05d30fc",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Inspect File Structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93cf52b3-65b6-4b38-bc3b-f392b3ef4b54",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "# --- 1. Create a sample JSON file for demonstration ---\n",
    "# In your actual use case, you would already have this file.\n",
    "files = ['../data/train_fraud_labels.json', '../data/mcc_codes.json']\n",
    "\n",
    "# --- 2. Load the JSON files and inspect their structure ---\n",
    "for file_path in files:\n",
    "    try:\n",
    "        # Open the file and load the JSON data into a Python object\n",
    "        with open(file_path, 'r') as f:\n",
    "            data = json.load(f)\n",
    "\n",
    "        print(\"Successfully loaded JSON file into memory.\")\n",
    "        print(f\"\\n--- Inspecting Object Structure for {file_path} ---\")\n",
    "\n",
    "        # Check the type of the top-level object\n",
    "        print(f\"1. Type of the loaded object: {type(data)}\")\n",
    "\n",
    "        # Check the top-level keys\n",
    "        if isinstance(data, dict):\n",
    "            top_level_keys = list(data.keys())\n",
    "            print(f\"2. Top-level keys: {top_level_keys}\")\n",
    "\n",
    "            # Inspect the nested object under the 'target' key\n",
    "            if 'target' in top_level_keys:\n",
    "                target_object = data['target']\n",
    "                print(\"\\n3. Inspecting the nested object under the 'target' key:\")\n",
    "                print(f\"  - Type of nested object: {type(target_object)}\")\n",
    "\n",
    "                if isinstance(target_object, dict):\n",
    "                    num_items = len(target_object)\n",
    "                    print(f\"  - Number of items (transaction ID/label pairs): {num_items}\")\n",
    "\n",
    "                    # Show a few example items without printing the whole dictionary\n",
    "                    print(\"  - Example items (first 5):\")\n",
    "                    count = 0\n",
    "                    for key, value in target_object.items():\n",
    "                        print(f\"    '{key}': '{value}'\")\n",
    "                        count += 1\n",
    "                        if count >= 5:\n",
    "                            break\n",
    "        else:\n",
    "            print(\"The JSON file does not contain a top-level dictionary as expected.\")\n",
    "\n",
    "\n",
    "    except FileNotFoundError:\n",
    "        print(f\"Error: The file '{file_path}' was not found.\")\n",
    "    except json.JSONDecodeError:\n",
    "        print(f\"Error: The file '{file_path}' is not a valid JSON file.\")\n",
    "    finally:\n",
    "        print(\"\\n--- Inspection Complete ---\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b478e275-56bd-47d9-adf1-119de83fce2d",
   "metadata": {},
   "source": [
    "#### Transform Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "773fa529-3d68-4b1e-98f4-96189615427c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import csv\n",
    "import os\n",
    "\n",
    "def convert_json_to_csv(json_path, csv_path, processor_func):\n",
    "    \"\"\"\n",
    "    A generic function to convert a JSON file to CSV using a specific processor.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        print(f\"Processing {json_path}...\")\n",
    "\n",
    "        # Step 1: Read the entire JSON file into memory and automatically close it.\n",
    "        with open(json_path, 'r') as json_file:\n",
    "            data = json.load(json_file)\n",
    "\n",
    "        # Step 2: Use the provided processor function to get the header and rows.\n",
    "        header, rows = processor_func(data)\n",
    "        \n",
    "        # Step 3: Open the destination CSV file and write the data.\n",
    "        with open(csv_path, 'w', newline='') as csv_file:\n",
    "            writer = csv.writer(csv_file)\n",
    "            writer.writerow(header)\n",
    "            writer.writerows(rows)\n",
    "        \n",
    "        print(f\"Successfully created {csv_path}\")\n",
    "\n",
    "    except FileNotFoundError:\n",
    "        print(f\"ERROR: File not found at {json_path}\")\n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred while processing {json_path}: {e}\")\n",
    "\n",
    "\n",
    "def process_fraud_labels(data):\n",
    "    \"\"\"Processor for the fraud labels JSON structure.\"\"\"\n",
    "    header = ['transaction_id', 'is_fraud']\n",
    "    # The actual data is nested under the 'target' key\n",
    "    target_data = data.get('target', {})\n",
    "    rows = list(target_data.items())\n",
    "    return header, rows\n",
    "\n",
    "\n",
    "def process_mcc_codes(data):\n",
    "    \"\"\"Processor for the MCC codes JSON structure.\"\"\"\n",
    "    header = ['mcc', 'description']\n",
    "    rows = list(data.items())\n",
    "    return header, rows\n",
    "\n",
    "\n",
    "# --- Main execution ---\n",
    "def main():\n",
    "    # Define file paths\n",
    "    data_directory = '../data'\n",
    "    \n",
    "    # Paths for the first conversion\n",
    "    fraud_json_path = os.path.join(data_directory, 'train_fraud_labels.json')\n",
    "    fraud_csv_path = os.path.join(data_directory, 'fraud_labels.csv') # Output file\n",
    "\n",
    "    # Paths for the second conversion\n",
    "    mcc_json_path = os.path.join(data_directory, 'mcc_codes.json')\n",
    "    mcc_csv_path = os.path.join(data_directory, 'mcc_codes.csv') # Output file\n",
    "    \n",
    "    # 3. Run the conversions using the generic function and specific processors\n",
    "    convert_json_to_csv(fraud_json_path, fraud_csv_path, process_fraud_labels)\n",
    "    print() # Add a newline for better readability\n",
    "    convert_json_to_csv(mcc_json_path, mcc_csv_path, process_mcc_codes)\n",
    "\n",
    "    print(\"\\n--- Conversion complete. ---\")\n",
    "    \n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "188ad959-b43e-4f6e-a075-c0dcccfdb69f",
   "metadata": {},
   "source": [
    "#### Inspect CSV Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bac944df-7017-4dc1-880e-f6a585012210",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!head -n 5 ../data/fraud_labels.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e7f9e0f-e701-49a1-9eda-80ea39682327",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!head -n 5 ../data/mcc_codes.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "546ab199-1792-4fc7-8b4b-d68b3999b4d5",
   "metadata": {},
   "source": [
    "#### Delete JSON Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57694100-1af3-4f41-87bd-3097ae800cad",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Remove json files after validating new csv files are good\n",
    "files_to_remove = ['../data/train_fraud_labels.json', '../data/mcc_codes.json']\n",
    "for file in files_to_remove:\n",
    "    os.remove(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1420b5d-ceac-4828-98ed-ef84690fc5c1",
   "metadata": {},
   "source": [
    "### Remove Header Rows\n",
    "\n",
    "This is required for the AlloyDB import later in the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad75bc1f-aa82-4dd0-8abc-32799172c16c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from tempfile import NamedTemporaryFile\n",
    "from shutil import move\n",
    "\n",
    "def remove_first_line(filepath):\n",
    "    \"\"\"\n",
    "    Removes the first line from a given file by writing all but the first\n",
    "    line to a temporary file and then replacing the original file.\n",
    "    \"\"\"\n",
    "    # Create a temporary file in the same directory to ensure it's on the same filesystem\n",
    "    # which allows for an atomic move operation.\n",
    "    file_dir = os.path.dirname(filepath)\n",
    "    try:\n",
    "        with open(filepath, 'r', encoding='utf-8') as f_in, NamedTemporaryFile(mode='w', encoding='utf-8', delete=False, dir=file_dir, suffix='.tmp') as f_out:\n",
    "            \n",
    "            temp_path = f_out.name\n",
    "            \n",
    "            # Skip the first line of the input file\n",
    "            next(f_in)\n",
    "            \n",
    "            # Copy the rest of the lines to the temporary file\n",
    "            for line in f_in:\n",
    "                f_out.write(line)\n",
    "\n",
    "        # Replace the original file with the temporary file\n",
    "        move(temp_path, filepath)\n",
    "        print(f\"Successfully removed the first line from: {filepath}\")\n",
    "\n",
    "    except StopIteration:\n",
    "        # This handles the case where the file is empty\n",
    "        print(f\"File is empty, no lines to remove: {filepath}\")\n",
    "        if 'temp_path' in locals() and os.path.exists(temp_path):\n",
    "            os.remove(temp_path)\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(f\"An error occurred while processing {filepath}: {e}\")\n",
    "        # Clean up the temporary file if an error occurs\n",
    "        if 'temp_path' in locals() and os.path.exists(temp_path):\n",
    "            os.remove(temp_path)\n",
    "\n",
    "def process_all_csv_in_directory(directory_path):\n",
    "    \"\"\"\n",
    "    Scans a directory for .csv files and removes the first line from each.\n",
    "    \"\"\"\n",
    "    # Check if the directory exists\n",
    "    if not os.path.isdir(directory_path):\n",
    "        print(f\"Error: Directory not found at '{directory_path}'\")\n",
    "        return\n",
    "\n",
    "    print(f\"Scanning for CSV files in '{directory_path}'...\")\n",
    "    \n",
    "    # Iterate over all files in the given directory\n",
    "    for filename in os.listdir(directory_path):\n",
    "        # Check if the file is a CSV file\n",
    "        if filename.lower().endswith('.csv'):\n",
    "            # Construct the full file path\n",
    "            file_path = os.path.join(directory_path, filename)\n",
    "            \n",
    "            # Call the function to remove the first line\n",
    "            remove_first_line(file_path)\n",
    "\n",
    "# --- Main execution ---\n",
    "def main():\n",
    "    # Specify the directory containing the CSV files\n",
    "    data_directory = '../data/'\n",
    "    \n",
    "    # Run the process\n",
    "    process_all_csv_in_directory(data_directory)\n",
    "    \n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36d7b975-35fc-4b5b-925d-f9ef43bcbceb",
   "metadata": {},
   "source": [
    "### Clean Transactions and Cards Data\n",
    "\n",
    "Some of the formatting in the transactions and cards data causes errors on import to AlloyDB. This section cleans up the data so that it can be imported successfully."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a301159a-b914-4ff5-9c9c-d1f675c539e8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "# --- Step 1: Clean the transactions data ---\n",
    "\n",
    "print(\"Cleaning transactions data...\")\n",
    "\n",
    "# Define column names based on your table structure\n",
    "transaction_cols = [\n",
    "    'id', 'date', 'client_id', 'card_id', 'amount', 'use_chip',\n",
    "    'merchant_id', 'merchant_city', 'merchant_state', 'zip', 'mcc', 'errors'\n",
    "]\n",
    "\n",
    "# Read the CSV file into a pandas DataFrame, assuming no header row\n",
    "try:\n",
    "    df_trans = pd.read_csv('../data/transactions_data.csv', header=None, names=transaction_cols)\n",
    "\n",
    "    # Clean the 'amount' column: remove '$' and convert to numeric\n",
    "    # The .loc accessor is used to safely modify the DataFrame\n",
    "    df_trans.loc[:, 'amount'] = df_trans['amount'].astype(str).str.replace('$', '').astype(float)\n",
    "\n",
    "    # Clean the 'zip' column: convert float (e.g., 58523.0) to integer then string\n",
    "    # Handling potential missing values (NaN) before converting\n",
    "    df_trans.loc[:, 'zip'] = df_trans['zip'].dropna().astype(int).astype(str)\n",
    "\n",
    "    # Save the cleaned data to a new CSV file without headers or index\n",
    "    cleaned_transactions_file = '../data/transactions_data.csv'\n",
    "    df_trans.to_csv(cleaned_transactions_file, index=False, header=False)\n",
    "\n",
    "    print(f\"Successfully cleaned transactions data and saved to '{cleaned_transactions_file}'\")\n",
    "    print(\"\\nPreview of cleaned transactions data:\")\n",
    "    print(df_trans.head().to_string())\n",
    "\n",
    "except FileNotFoundError:\n",
    "    print(\"Error: 'transactions.csv' not found. Please ensure the file is in the correct directory.\")\n",
    "except Exception as e:\n",
    "    print(f\"An error occurred while processing transactions: {e}\")\n",
    "\n",
    "\n",
    "# --- Step 2: Clean the cards data ---\n",
    "\n",
    "print(\"\\n\" + \"=\"*50 + \"\\n\")\n",
    "print(\"Cleaning cards data...\")\n",
    "\n",
    "# Define column names based on your table structure\n",
    "card_cols = [\n",
    "    'id', 'client_id', 'card_brand', 'card_type', 'card_number', 'expires',\n",
    "    'cvv', 'has_chip', 'num_cards_issued', 'credit_limit', 'acct_open_date',\n",
    "    'year_pin_last_changed', 'card_on_dark_web'\n",
    "]\n",
    "\n",
    "# Read the CSV file, assuming no header row\n",
    "try:\n",
    "    df_cards = pd.read_csv('../data/cards_data.csv', header=None, names=card_cols)\n",
    "\n",
    "    # Clean 'acct_open_date': convert 'MM/YYYY' to 'YYYY-MM-01'\n",
    "    df_cards.loc[:, 'acct_open_date'] = pd.to_datetime(df_cards['acct_open_date'], format='%m/%Y').dt.strftime('%Y-%m-01')\n",
    "\n",
    "    # Clean 'credit_limit': remove '$' (PostgreSQL's MONEY type can be picky)\n",
    "    df_cards.loc[:, 'credit_limit'] = df_cards['credit_limit'].astype(str).str.replace('$', '').astype(float)\n",
    "\n",
    "    # Save the cleaned data to a new CSV file\n",
    "    cleaned_cards_file = '../data/cards_data.csv'\n",
    "    df_cards.to_csv(cleaned_cards_file, index=False, header=False)\n",
    "\n",
    "    print(f\"Successfully cleaned cards data and saved to '{cleaned_cards_file}'\")\n",
    "    print(\"\\nPreview of cleaned cards data:\")\n",
    "    print(df_cards.head().to_string())\n",
    "\n",
    "except FileNotFoundError:\n",
    "    print(\"Error: 'cards.csv' not found. Please ensure the file is in the correct directory.\")\n",
    "except Exception as e:\n",
    "    print(f\"An error occurred while processing cards: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf5fbec7-97d5-47fa-8296-ec95c358e2d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!head -n 5 ../data/transactions_data.csv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78413465-1952-4e3d-9969-3db8eb378bc4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!head -n 5 ../data/cards_data.csv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26a36336-9a7a-4ad4-8faa-399ad1348093",
   "metadata": {},
   "source": [
    "### Upload Files to GCS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81d6c1a1-662f-4a0a-b642-d11b747b2b33",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from google.cloud import storage\n",
    "\n",
    "def upload_directory_to_gcs(bucket_name, source_directory, destination_blob_prefix):\n",
    "    \"\"\"Uploads all files from a directory to a GCS bucket.\"\"\"\n",
    "\n",
    "    # Ensure the source directory exists\n",
    "    if not os.path.isdir(source_directory):\n",
    "        print(f\"Error: The source directory '{source_directory}' does not exist.\")\n",
    "        return\n",
    "\n",
    "    # Initialize the Google Cloud Storage client\n",
    "    storage_client = storage.Client()\n",
    "    bucket = storage_client.bucket(bucket_name)\n",
    "\n",
    "    print(f\"Uploading files from '{source_directory}' to 'gs://{bucket_name}/{destination_blob_prefix}'...\")\n",
    "\n",
    "    # Walk through the source directory\n",
    "    for dirpath, _, filenames in os.walk(source_directory):\n",
    "        for filename in filenames:\n",
    "            # Construct the full local path of the file\n",
    "            local_file_path = os.path.join(dirpath, filename)\n",
    "\n",
    "            # Construct the destination path in GCS\n",
    "            relative_path = os.path.relpath(local_file_path, source_directory)\n",
    "            destination_blob_name = os.path.join(destination_blob_prefix, relative_path)\n",
    "\n",
    "            # Create a blob and upload the file\n",
    "            blob = bucket.blob(destination_blob_name)\n",
    "            blob.upload_from_filename(local_file_path)\n",
    "\n",
    "            print(f\"  Uploaded {local_file_path} to gs://{bucket_name}/{destination_blob_name}\")\n",
    "\n",
    "    print(\"Upload complete.\")\n",
    "\n",
    "def main():\n",
    "    # The local directory you want to upload\n",
    "    local_directory = \"../data/\"\n",
    "\n",
    "    # The destination directory within your GCS bucket\n",
    "    gcs_destination_directory = \"data/\"\n",
    "\n",
    "    upload_directory_to_gcs(gcs_bucket_name, local_directory, gcs_destination_directory)\n",
    "    \n",
    "main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "548f183a-de8d-4b21-8e12-1c7ec26bcdfc",
   "metadata": {},
   "source": [
    "## Setup Database"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "000dccda-3c9d-480e-9715-e990c38632da",
   "metadata": {},
   "source": [
    "### Connect to the AlloyDB Cluster\n",
    "\n",
    "This function will create a connection pool to your AlloyDB instance using the AlloyDB Python connector. The AlloyDB Python connector will automatically create secure connections to your AlloyDB instance using mTLS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53e8eea3-deac-474e-b6c3-e03d1b3d49c5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import asyncpg\n",
    "\n",
    "import sqlalchemy\n",
    "from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine\n",
    "\n",
    "from google.cloud.alloydb.connector import AsyncConnector, IPTypes\n",
    "\n",
    "async def init_connection_pool(connector: AsyncConnector, db_name: str = alloydb_database, pool_size: int = 5) -> AsyncEngine:\n",
    "    # initialize Connector object for connections to AlloyDB\n",
    "    connection_string = f\"projects/{project_id}/locations/{region}/clusters/{alloydb_cluster}/instances/{alloydb_instance}\"\n",
    "\n",
    "    async def getconn() -> asyncpg.Connection:\n",
    "        conn: asyncpg.Connection = await connector.connect(\n",
    "            connection_string,\n",
    "            \"asyncpg\",\n",
    "            user=\"postgres\",\n",
    "            password=alloydb_password,\n",
    "            db=db_name,\n",
    "            ip_type=IPTypes.PRIVATE, # Optionally use IPTypes.PUBLIC\n",
    "        )\n",
    "        return conn\n",
    "\n",
    "    pool = create_async_engine(\n",
    "        \"postgresql+asyncpg://\",\n",
    "        async_creator=getconn,\n",
    "        pool_size=pool_size,\n",
    "        max_overflow=0,\n",
    "        isolation_level='AUTOCOMMIT'\n",
    "    )\n",
    "    return pool\n",
    "\n",
    "connector = AsyncConnector()\n",
    "\n",
    "postgres_db_pool = await init_connection_pool(connector, \"postgres\") # Need to connect to postgres database first to create the application database\n",
    "finance_db_pool = await init_connection_pool(connector, f\"{alloydb_database}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf3fada4-471d-4135-92bf-6dbfbdb48e56",
   "metadata": {},
   "source": [
    "### Create the Finance Database"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac2ae811-2137-453e-ad60-27225d892a6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Close existing connections to the database\n",
    "sql = f\"\"\"SELECT pg_terminate_backend(pg_stat_activity.pid)\n",
    "FROM pg_stat_activity\n",
    "WHERE pg_stat_activity.datname = '{alloydb_database}'\n",
    "  AND pid <> pg_backend_pid();\"\"\"\n",
    "await run_query(postgres_db_pool, sql)\n",
    "\n",
    "# Uncomment this section and run this cell to drop the existing database if you want to replace it.\n",
    "# Drop the database\n",
    "#sql = f\"DROP DATABASE {alloydb_database};\"\n",
    "#await run_query(postgres_db_pool, sql)\n",
    "\n",
    "# Reinitiate the connection pool\n",
    "finance_db_pool = await init_connection_pool(connector, f\"{alloydb_database}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "520d568b-ef56-4096-88af-088b7a7482cb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Use postgres_db_pool to create the database\n",
    "sql = f\"CREATE DATABASE {alloydb_database};\"\n",
    "result = await run_query(postgres_db_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "670e07b0-4383-44b1-9aee-e413b15495ac",
   "metadata": {},
   "source": [
    "### Create Tables and Extensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4347491e-339e-44cb-85a5-e38fad7f61bb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Use finance_db_pool to create the rest of the schema objects\n",
    "\n",
    "sql_array = []\n",
    "\n",
    "sql_array.append(\"CREATE EXTENSION IF NOT EXISTS vector;\")\n",
    "\n",
    "sql_array.append(\"CREATE EXTENSION IF NOT EXISTS google_ml_integration;\")\n",
    "\n",
    "sql_array.append(\"\"\"CREATE TABLE transactions (\n",
    "    id BIGINT PRIMARY KEY,\n",
    "    date TIMESTAMP WITHOUT TIME ZONE,\n",
    "    client_id INTEGER,\n",
    "    card_id INTEGER,\n",
    "    amount NUMERIC(10, 2),\n",
    "    use_chip VARCHAR(255),\n",
    "    merchant_id INTEGER,\n",
    "    merchant_city VARCHAR(255),\n",
    "    merchant_state VARCHAR(50),\n",
    "    zip VARCHAR(10),\n",
    "    mcc INTEGER,\n",
    "    errors TEXT\n",
    ");\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"CREATE TABLE cards (\n",
    "    id BIGINT PRIMARY KEY,\n",
    "    client_id INTEGER,\n",
    "    card_brand VARCHAR(50),\n",
    "    card_type VARCHAR(50),\n",
    "    card_number VARCHAR(20),\n",
    "    expires VARCHAR(7),\n",
    "    cvv VARCHAR(4),\n",
    "    has_chip BOOLEAN,\n",
    "    num_cards_issued INTEGER,\n",
    "    credit_limit MONEY,\n",
    "    acct_open_date DATE,\n",
    "    year_pin_last_changed INTEGER,\n",
    "    card_on_dark_web BOOLEAN\n",
    ");\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"CREATE TABLE users (\n",
    "    id BIGINT PRIMARY KEY,\n",
    "    current_age INTEGER,\n",
    "    retirement_age INTEGER,\n",
    "    birth_year INTEGER,\n",
    "    birth_month INTEGER,\n",
    "    gender VARCHAR(10),\n",
    "    address VARCHAR(255),\n",
    "    latitude NUMERIC(10, 8),\n",
    "    longitude NUMERIC(11, 8),\n",
    "    per_capita_income MONEY,\n",
    "    yearly_income MONEY,\n",
    "    total_debt MONEY,\n",
    "    credit_score INTEGER,\n",
    "    num_credit_cards INTEGER\n",
    ");\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"CREATE TABLE mcc_codes (\n",
    "    mcc SMALLINT PRIMARY KEY,\n",
    "    description VARCHAR(255) NOT NULL\n",
    ");\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"CREATE TABLE fraud_labels (\n",
    "    transaction_id BIGINT PRIMARY KEY,\n",
    "    is_fraud BOOLEAN NOT NULL\n",
    ");\"\"\")\n",
    "\n",
    "for sql in sql_array:\n",
    "  await run_query(finance_db_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8ceb01-5981-48da-8b15-941d4426cbba",
   "metadata": {},
   "source": [
    "### Import Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d0c577c-ac6a-49aa-9776-52c8c59a0179",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### Run the Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcbfc5c0-fc94-4dd0-ba89-e7582cd33437",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Reference: https://cloud.google.com/alloydb/docs/import-csv-file#rest-v1\n",
    "#            https://cloud.google.com/alloydb/docs/reference/rest/v1/projects.locations.operations/get\n",
    "\n",
    "import time\n",
    "\n",
    "# Define tables to import\n",
    "# Format: [ [ table, [list-of-files] ] ]\n",
    "import_array = [\n",
    "    ['users',['users_data.csv']],\n",
    "    ['mcc_codes',['mcc_codes.csv']],\n",
    "    ['transactions',['transactions_data.csv']],\n",
    "    ['fraud_labels',['fraud_labels.csv']],\n",
    "    ['cards',['cards_data.csv']]\n",
    "]\n",
    "\n",
    "def import_csv_to_alloydb(import_array):\n",
    "\n",
    "  operations = []\n",
    "\n",
    "  for table, files in import_array:\n",
    "    for f in files:\n",
    "      url = f\"https://alloydb.googleapis.com/v1/projects/{project_id}/locations/{region}/clusters/{alloydb_cluster}:import\"\n",
    "      request_body = {\n",
    "        \"gcsUri\": f\"gs://{gcs_bucket_name}/data/{f}\",\n",
    "        \"database\": f\"{alloydb_database}\",\n",
    "        \"user\": \"postgres\",\n",
    "        \"csvImportOptions\": {\n",
    "          \"table\": f\"{table}\",\n",
    "          #\"columns\": [\"COLUMN1\", \"COLUMN2\"],\n",
    "          #\"fieldDelimiter\": \"FIELD_DELIMITER\",\n",
    "          #\"quoteCharacter\": \"QUOTE_CHARACTER\",\n",
    "          #\"escapeCharacter\": \"ESCAPE_CHARACTER\"\n",
    "        }\n",
    "      }\n",
    "      response = rest_api_helper(authed_session, url, 'POST', request_body, {})\n",
    "      operations.append([table, response['name']])\n",
    "      print(response)\n",
    "\n",
    "  for o in operations:\n",
    "    operation_complete = False\n",
    "    while operation_complete == False:\n",
    "      print(f\"Operation for table {o[0]} still running: {o[1]}\")\n",
    "      url = f\"https://alloydb.googleapis.com/v1/{o[1]}\"\n",
    "      response = rest_api_helper(authed_session, url, 'GET', request_body, {})\n",
    "      operation_complete = response['done']\n",
    "      if operation_complete:\n",
    "        print(f\"Import complete for table: {o[0]}. \\nResult: {response}\")\n",
    "        continue\n",
    "      time.sleep(5)\n",
    "\n",
    "  return \"Import operation complete.\"\n",
    "\n",
    "import_csv_to_alloydb(import_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "911ec780-20f5-4434-b390-24cc0ebc7023",
   "metadata": {},
   "source": [
    "#### Validate Row Counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d0b0d86-cabf-4f74-a435-effde2dac65f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql = \"\"\"\n",
    "SELECT 'users' AS table_name, (SELECT COUNT(*) FROM users) AS imported_count, 2000 AS target_row_count\n",
    "UNION ALL\n",
    "SELECT 'mcc_codes', (SELECT COUNT(*) FROM mcc_codes), 109\n",
    "UNION ALL\n",
    "SELECT 'transactions', (SELECT COUNT(*) FROM transactions), 13305915\n",
    "UNION ALL\n",
    "SELECT 'fraud_labels', (SELECT COUNT(*) FROM fraud_labels), 8914963\n",
    "UNION ALL\n",
    "SELECT 'cards', (SELECT COUNT(*) FROM cards), 6146;\n",
    "\"\"\"\n",
    "\n",
    "await run_query(finance_db_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b6ada19-2d3d-46e1-a374-9b873894d369",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Setup Referential Integrity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "293afaa0-44fa-4410-9a8c-66888a7ebf43",
   "metadata": {},
   "source": [
    "#### Create Foreign Keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57830316-eccc-49c2-914e-53f1c88bfb6f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql_array = []\n",
    "\n",
    "sql_array.append(\"\"\"-- In the 'transactions' table, link 'client_id' to the 'users' table.\n",
    "ALTER TABLE transactions\n",
    "ADD CONSTRAINT fk_transactions_client\n",
    "FOREIGN KEY (client_id) REFERENCES users(id);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- In the 'transactions' table, link 'card_id' to the 'cards' table.\n",
    "ALTER TABLE transactions\n",
    "ADD CONSTRAINT fk_transactions_card\n",
    "FOREIGN KEY (card_id) REFERENCES cards(id);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- In the 'transactions' table, link 'mcc' to the 'mcc_codes' table.\n",
    "ALTER TABLE transactions\n",
    "ADD CONSTRAINT fk_transactions_mcc\n",
    "FOREIGN KEY (mcc) REFERENCES mcc_codes(mcc);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- In the 'cards' table, link 'client_id' to the 'users' table.\n",
    "ALTER TABLE cards\n",
    "ADD CONSTRAINT fk_cards_client\n",
    "FOREIGN KEY (client_id) REFERENCES users(id);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- In the 'fraud_labels' table, link 'transaction_id' to the 'transactions' table.\n",
    "ALTER TABLE fraud_labels\n",
    "ADD CONSTRAINT fk_fraud_labels_transaction\n",
    "FOREIGN KEY (transaction_id) REFERENCES transactions(id);\n",
    "\"\"\")\n",
    "\n",
    "for sql in sql_array:\n",
    "  await run_query(finance_db_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5450c91a-4dbe-4ccb-9248-95725f461cd4",
   "metadata": {},
   "source": [
    "#### Validate Foreign Keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b2e5e8b-e71e-4ed9-bfba-a83ef02d13ed",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql_array = []\n",
    "\n",
    "# Validate FK constraints\n",
    "sql_array.append(\"ALTER TABLE transactions VALIDATE CONSTRAINT fk_transactions_client;\")\n",
    "sql_array.append(\"ALTER TABLE transactions VALIDATE CONSTRAINT fk_transactions_card;\")\n",
    "sql_array.append(\"ALTER TABLE transactions VALIDATE CONSTRAINT fk_transactions_mcc;\")\n",
    "sql_array.append(\"ALTER TABLE cards VALIDATE CONSTRAINT fk_cards_client;\")\n",
    "sql_array.append(\"ALTER TABLE fraud_labels VALIDATE CONSTRAINT fk_fraud_labels_transaction;\")\n",
    "\n",
    "for sql in sql_array:\n",
    "  response = await run_query(finance_db_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21ef2720-54f1-45e5-a39b-298668f81818",
   "metadata": {},
   "source": [
    "### Create Indexes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a931fd79-2a9d-4c2a-b333-d32360d95ec1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql_array = []\n",
    "\n",
    "# FOREIGN KEY INDEXES\n",
    "\n",
    "sql_array.append(\"\"\"-- Index on transactions.client_id to speed up joins to the users table\n",
    "CREATE INDEX idx_transactions_client_id ON transactions (client_id);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Index on transactions.card_id to speed up joins to the cards table\n",
    "CREATE INDEX idx_transactions_card_id ON transactions (card_id);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Index on transactions.mcc to speed up joins to the mcc_codes table\n",
    "CREATE INDEX idx_transactions_mcc ON transactions (mcc);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Index on cards.client_id to speed up finding all cards for a given user\n",
    "CREATE INDEX idx_cards_client_id ON cards (client_id);\n",
    "\"\"\")\n",
    "\n",
    "# INDEXES FOR COMMON QUERIES \n",
    "\n",
    "sql_array.append(\"\"\"-- Index on the transaction date for time-series analysis\n",
    "CREATE INDEX idx_transactions_date ON transactions (date);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Index on merchant_id for quickly finding transactions for a specific merchant\n",
    "CREATE INDEX idx_transactions_merchant_id ON transactions (merchant_id);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Composite index on merchant location for geographic queries\n",
    "CREATE INDEX idx_transactions_merchant_location ON transactions (merchant_state, merchant_city);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Unique index on card_number for direct card lookups\n",
    "CREATE UNIQUE INDEX idx_cards_card_number_unique ON cards (card_number);\n",
    "\"\"\")\n",
    "\n",
    "sql_array.append(\"\"\"-- Index on credit_score in the users table for analytical queries\n",
    "CREATE INDEX idx_users_credit_score ON users (credit_score);\n",
    "\"\"\")\n",
    "\n",
    "# SPECIALIZED PARTIAL INDEX\n",
    "\n",
    "sql_array.append(\"\"\"-- Highly efficient index that ONLY includes fraudulent transactions\n",
    "CREATE INDEX idx_fraud_labels_is_fraud_true ON fraud_labels (transaction_id) WHERE is_fraud = TRUE;\n",
    "\"\"\")\n",
    "\n",
    "for sql in sql_array:\n",
    "  response = await run_query(finance_db_pool, sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c2ef46d-e9b2-4349-aa93-b5ed40a6d65f",
   "metadata": {},
   "source": [
    "## Export to SQL File\n",
    "\n",
    "You can now export the database to a .sql file to speed up and simplify future data loads for this dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "decc3ca0-e07c-4af0-916c-5d07ce5bd758",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "url = f\"https://alloydb.googleapis.com/v1/projects/{project_id}/locations/{region}/clusters/{alloydb_cluster}:export\"\n",
    "request_body = {\n",
    "  \"gcs_destination\": {\n",
    "    \"uri\": f\"gs://{gcs_bucket_name}/export/{alloydb_database}.sql\"\n",
    "  },\n",
    "  \"database\": f\"{alloydb_database}\",\n",
    "  \"sql_export_options\": {\n",
    "    \"schema_only\": \"false\",\n",
    "    \"tables\": [\n",
    "        'users',\n",
    "        'mcc_codes',\n",
    "        'transactions',\n",
    "        'fraud_labels',\n",
    "        'cards'\n",
    "    ],\n",
    "    \"clean_target_objects\": \"false\",\n",
    "    \"if_exist_target_objects\": \"false\"\n",
    "  }\n",
    "}\n",
    "\n",
    "result = rest_api_helper(authed_session, url, 'POST', request_body, {})\n",
    "print(f\"Kicked off export: {result}\")\n",
    "\n",
    "operation_id = result['name']\n",
    "\n",
    "operation_complete = False\n",
    "while operation_complete == False:\n",
    "  print(f\"Export still running: {operation_id}\")\n",
    "  url = f\"https://alloydb.googleapis.com/v1/{operation_id}\"\n",
    "  response = rest_api_helper(authed_session, url, 'GET', request_body, {})\n",
    "  operation_complete = response['done']\n",
    "  if operation_complete:\n",
    "    print(f\"Operation complete. \\nResult: {response}\")\n",
    "    continue\n",
    "  time.sleep(5)"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-base-py",
   "name": "workbench-notebooks.m130",
   "type": "gcloud",
   "uri": "us-docker.pkg.dev/deeplearning-platform-release/gcr.io/workbench-notebooks:m130"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel) (Local)",
   "language": "python",
   "name": "conda-base-py"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
